# CR模块 - 屏幕内容理解
import os
import logging
import sys
import time
from PIL import Image
from typing import Dict, Any, List, Optional

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("CR_Module")

# 尝试导入pyautogui
pyautogui = None
try:
    import pyautogui
    logger.info("pyautogui导入成功")
except ImportError:
    logger.warning("pyautogui未安装，截图功能将不可用")
    logger.warning("请运行: pip install pyautogui")

# 尝试导入numpy
np = None
try:
    import numpy as np
    logger.info("numpy导入成功")
except ImportError:
    logger.warning("numpy未安装，图像处理功能将受限")
    logger.warning("请运行: pip install numpy")

# 尝试导入OpenCV
cv2 = None
try:
    import cv2
    logger.info("OpenCV导入成功")
except ImportError:
    logger.warning("OpenCV未安装，图像处理功能将受限")
    logger.warning("请运行: pip install opencv-python")

# 尝试导入pytesseract
pytesseract = None
tesseract_available = False
try:
    import pytesseract
    # 尝试设置tesseract_cmd，如果环境变量未设置
    try:
        # 尝试获取tesseract路径
        if not hasattr(pytesseract, 'tesseract_cmd') or not pytesseract.tesseract_cmd:
            # 常见的Tesseract安装路径
            common_paths = [
                r'C:\Program Files\Tesseract-OCR\tesseract.exe',
                r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
                r'D:\Program Files\Tesseract-OCR\tesseract.exe',
                r'D:\Program Files (x86)\Tesseract-OCR\tesseract.exe'
            ]
            
            # 搜索系统环境变量PATH
            for path in os.environ.get('PATH', '').split(';'):
                if path and os.path.exists(os.path.join(path, 'tesseract.exe')):
                    common_paths.insert(0, os.path.join(path, 'tesseract.exe'))
                    break
            
            # 尝试设置有效的tesseract路径
            for tesseract_path in common_paths:
                if os.path.exists(tesseract_path):
                    pytesseract.pytesseract.tesseract_cmd = tesseract_path
                    logger.info(f"自动配置Tesseract路径: {tesseract_path}")
                    break
            
            # 如果没有找到，记录提示信息
            if not hasattr(pytesseract, 'tesseract_cmd') or not pytesseract.tesseract_cmd:
                logger.warning("未找到Tesseract OCR可执行文件")
        
        # 测试tesseract是否可用
        if hasattr(pytesseract, 'tesseract_cmd') and pytesseract.tesseract_cmd and os.path.exists(pytesseract.tesseract_cmd):
            try:
                # 使用Tesseract版本检查或简单测试
                if hasattr(pytesseract, 'get_tesseract_version'):
                    pytesseract.get_tesseract_version()
                elif np is not None:
                    # 使用一个简单的测试图像
                    test_img = np.zeros((10, 10, 3), dtype=np.uint8)
                    pytesseract.image_to_string(test_img)
                tesseract_available = True
                logger.info("Tesseract OCR可用")
            except Exception as e:
                tesseract_available = False
                logger.warning(f"Tesseract OCR测试失败: {str(e)}")
        else:
            tesseract_available = False
            logger.warning("Tesseract OCR路径未设置")
    except Exception as e:
        tesseract_available = False
        logger.warning(f"Tesseract OCR配置失败: {str(e)}")
except ImportError:
    logger.warning("pytesseract未安装，文本识别功能将不可用")
    logger.warning("请运行: pip install pytesseract")

# 定义一个函数来检查和设置Tesseract路径
def set_tesseract_path(path: str) -> bool:
    """
    手动设置Tesseract OCR的路径
    
    Args:
        path: Tesseract可执行文件的路径
    
    Returns:
        bool: 是否设置成功
    """
    global tesseract_available
    if pytesseract is None:
        logger.error("pytesseract未安装")
        return False
    
    if not os.path.exists(path):
        logger.error(f"指定的Tesseract路径不存在: {path}")
        return False
    
    try:
        pytesseract.pytesseract.tesseract_cmd = path
        # 测试路径是否有效
        if hasattr(pytesseract, 'get_tesseract_version'):
            pytesseract.get_tesseract_version()
        elif np is not None:
            test_img = np.zeros((10, 10, 3), dtype=np.uint8)
            pytesseract.image_to_string(test_img)
        tesseract_available = True
        logger.info(f"成功设置Tesseract路径: {path}")
        return True
    except Exception as e:
        logger.error(f"设置Tesseract路径失败: {str(e)}")
        tesseract_available = False
        return False

# 定义一个函数来自动搜索Tesseract安装路径
def find_tesseract() -> Optional[str]:
    """
    自动搜索系统中可能的Tesseract安装路径
    
    Returns:
        Optional[str]: 找到的Tesseract路径，如果没有找到则返回None
    """
    # 检查常见的安装位置
    common_paths = [
        r'C:\Program Files\Tesseract-OCR\tesseract.exe',
        r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
        r'D:\Program Files\Tesseract-OCR\tesseract.exe',
        r'D:\Program Files (x86)\Tesseract-OCR\tesseract.exe'
    ]
    
    # 检查环境变量PATH
    for path in os.environ.get('PATH', '').split(';'):
        if path:
            tesseract_path = os.path.join(path, 'tesseract.exe')
            if os.path.exists(tesseract_path):
                logger.info(f"在PATH中找到Tesseract: {tesseract_path}")
                return tesseract_path
    
    # 检查常见路径
    for tesseract_path in common_paths:
        if os.path.exists(tesseract_path):
            logger.info(f"在常见位置找到Tesseract: {tesseract_path}")
            return tesseract_path
    
    logger.warning("未找到Tesseract OCR")
    return None

class ScreenContentRecognizer:
    """屏幕内容识别器"""
    
    def __init__(self):
        self.pyautogui_available = pyautogui is not None
        self.numpy_available = np is not None
        self.opencv_available = cv2 is not None
        self.tesseract_available = tesseract_available
        self.ocr_config = '--oem 3 --psm 6'
        
        # 尝试自动查找Tesseract（如果尚未找到）
        if not self.tesseract_available and pytesseract is not None:
            tesseract_path = find_tesseract()
            if tesseract_path:
                self.tesseract_available = set_tesseract_path(tesseract_path)
        
        logger.info(f"屏幕内容识别器初始化 - 功能状态: 屏幕截图={self.pyautogui_available}, " \
                   f"numpy={self.numpy_available}, OpenCV={self.opencv_available}, " \
                   f"Tesseract OCR={self.tesseract_available}")
        
        # 如果Tesseract不可用，提供安装指南
        if not self.tesseract_available:
            logger.warning("未找到可用的Tesseract OCR。要启用OCR功能，请:")
            logger.warning("1. 从 https://github.com/UB-Mannheim/tesseract/wiki 下载并安装Tesseract OCR")
            logger.warning("2. 将Tesseract添加到系统PATH环境变量")
            logger.warning("3. 或在代码中使用 set_tesseract_path() 函数手动设置路径")
    
    def capture_screen(self, region: tuple = None) -> Image.Image:
        """捕获屏幕图像
        
        Args:
            region: (left, top, width, height) 区域坐标，None表示全屏
            
        Returns:
            屏幕图像
        """
        try:
            screenshot = pyautogui.screenshot(region=region)
            logger.info(f"屏幕截图成功{' (区域模式)' if region else ' (全屏模式)'}")
            return screenshot
        except Exception as e:
            logger.error(f"屏幕截图失败: {str(e)}")
            raise
    
    def extract_text(self, image: Image.Image = None, region: tuple = None, lang: str = 'chi_sim+eng') -> str:
        """从屏幕或图像中提取文本
        
        Args:
            image: 要提取文本的图像，如果为None则截图
            region: 如果没有提供图像，此参数指定截图区域
            lang: OCR语言，默认为中英文
            
        Returns:
            提取的文本
        """
        if not self.tesseract_available:
            # 尝试再次查找Tesseract
            if pytesseract is not None:
                tesseract_path = find_tesseract()
                if tesseract_path:
                    self.tesseract_available = set_tesseract_path(tesseract_path)
                    
            if not self.tesseract_available:
                logger.error("Tesseract OCR不可用，无法提取文本")
                return "错误: Tesseract OCR未安装或配置错误，请先安装Tesseract"
        
        try:
            # 如果没有提供图像，则截图
            if image is None:
                image = self.capture_screen(region)
            
            # 图像预处理以提高OCR准确性
            processed_image = self._preprocess_image(image)
            
            # 执行OCR
            text = pytesseract.image_to_string(processed_image, lang=lang, config=self.ocr_config)
            logger.info(f"OCR文本提取成功，提取到{len(text.strip())}个字符")
            return text.strip()
        except Exception as e:
            logger.error(f"文本提取失败: {str(e)}")
            # 如果是TesseractNotFoundError，尝试提供更多帮助
            if "tesseract is not installed" in str(e).lower() or "could not find tesseract" in str(e).lower():
                logger.warning("请确保Tesseract OCR已正确安装并添加到系统PATH")
                logger.warning("下载地址: https://github.com/UB-Mannheim/tesseract/wiki")
            return f"错误: {str(e)}"
    
    def _preprocess_image(self, image: Image.Image) -> Image.Image:
        """预处理图像以提高OCR准确性
        
        Args:
            image: 原始图像
            
        Returns:
            预处理后的图像
        """
        if not self.opencv_available:
            # 如果没有OpenCV，返回原始图像
            return image
        
        try:
            # 转换为OpenCV格式
            img = np.array(image)
            img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            
            # 应用高斯模糊减少噪点
            img_blur = cv2.GaussianBlur(img_gray, (5, 5), 0)
            
            # 自适应阈值处理
            img_threshold = cv2.adaptiveThreshold(
                img_blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY, 11, 2
            )
            
            # 转换回PIL格式
            return Image.fromarray(img_threshold)
        except Exception as e:
            logger.warning(f"图像预处理失败，使用原始图像: {str(e)}")
            return image
    
    def detect_text_regions(self, image: Image.Image = None, region: tuple = None) -> List[Dict[str, Any]]:
        """检测文本区域
        
        Args:
            image: 要分析的图像，如果为None则截图
            region: 如果没有提供图像，此参数指定截图区域
            
        Returns:
            文本区域列表，每个区域包含文本和坐标
        """
        if not self.tesseract_available:
            logger.error("Tesseract OCR不可用，无法检测文本区域")
            return []
        
        try:
            # 如果没有提供图像，则截图
            if image is None:
                image = self.capture_screen(region)
            
            # 获取OCR数据
            data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
            
            # 处理识别结果
            regions = []
            n_boxes = len(data['level'])
            
            for i in range(n_boxes):
                if int(data['conf'][i]) > 60:  # 只保留置信度大于60的结果
                    text = data['text'][i].strip()
                    if text:  # 非空文本
                        regions.append({
                            'text': text,
                            'confidence': data['conf'][i],
                            'x': data['left'][i],
                            'y': data['top'][i],
                            'width': data['width'][i],
                            'height': data['height'][i]
                        })
            
            logger.info(f"成功检测到{len(regions)}个文本区域")
            return regions
        except Exception as e:
            logger.error(f"文本区域检测失败: {str(e)}")
            return []
    
    def analyze_screen_content(self, image: Image.Image = None) -> Dict[str, Any]:
        """综合分析屏幕内容
        
        Args:
            image: 要分析的图像，如果为None则截图
            
        Returns:
            包含屏幕分析结果的字典
        """
        try:
            # 如果没有提供图像，则截图
            if image is None:
                image = self.capture_screen()
            
            # 提取所有文本
            full_text = self.extract_text(image)
            
            # 检测文本区域
            text_regions = self.detect_text_regions(image)
            
            # 构建分析结果
            analysis = {
                'full_text': full_text,
                'text_regions': text_regions,
                'text_count': len(full_text.strip()),
                'region_count': len(text_regions),
                'timestamp': pyautogui.time.time()
            }
            
            logger.info(f"屏幕内容分析完成: 提取到{analysis['text_count']}个字符，{analysis['region_count']}个文本区域")
            return analysis
        except Exception as e:
            logger.error(f"屏幕内容分析失败: {str(e)}")
            return {'error': str(e)}
    
    def search_text_on_screen(self, keyword: str, region: tuple = None) -> List[Dict[str, Any]]:
        """在屏幕上搜索特定文本
        
        Args:
            keyword: 要搜索的关键词
            region: 搜索区域，None表示全屏
            
        Returns:
            匹配区域列表
        """
        try:
            # 检测所有文本区域
            regions = self.detect_text_regions(region=region)
            
            # 查找匹配的区域
            matches = []
            for region_data in regions:
                if keyword.lower() in region_data['text'].lower():
                    matches.append(region_data)
            
            logger.info(f"在屏幕上找到{len(matches)}个'{keyword}'的匹配项")
            return matches
        except Exception as e:
            logger.error(f"屏幕文本搜索失败: {str(e)}")
            return []
    
    def get_screenshot_info(self) -> Dict[str, Any]:
        """获取当前屏幕信息
        
        Returns:
            屏幕信息字典
        """
        try:
            # 获取屏幕尺寸
            screen_width, screen_height = pyautogui.size()
            
            # 获取鼠标位置
            mouse_x, mouse_y = pyautogui.position()
            
            return {
                'screen_size': (screen_width, screen_height),
                'mouse_position': (mouse_x, mouse_y),
                'tesseract_available': self.tesseract_available,
                'opencv_available': self.opencv_available
            }
        except Exception as e:
            logger.error(f"获取屏幕信息失败: {str(e)}")
            return {'error': str(e)}

# 创建全局实例
screen_recognizer = ScreenContentRecognizer()

# 导出常用函数
def extract_text_from_screen(region: tuple = None, lang: str = 'chi_sim+eng') -> str:
    """从屏幕提取文本的便捷函数"""
    return screen_recognizer.extract_text(region=region, lang=lang)

def detect_text_on_screen(region: tuple = None) -> List[Dict[str, Any]]:
    """检测屏幕文本区域的便捷函数"""
    return screen_recognizer.detect_text_regions(region=region)

def search_text(text: str, region: tuple = None) -> List[Dict[str, Any]]:
    """搜索屏幕文本的便捷函数"""
    return screen_recognizer.search_text_on_screen(text, region=region)

def analyze_screen() -> Dict[str, Any]:
    """分析屏幕内容的便捷函数"""
    return screen_recognizer.analyze_screen_content()

def get_screen_info() -> Dict[str, Any]:
    """获取屏幕信息的便捷函数"""
    return screen_recognizer.get_screenshot_info()

# 模块可用性检查
def is_available() -> bool:
    """检查屏幕内容识别模块是否可用
    
    Returns:
        bool: 模块是否可用（只要有任一依赖可用，模块就被视为可用）
    """
    # 只要有任一依赖可用，模块就被视为可用
    return pyautogui is not None or np is not None or cv2 is not None or tesseract_available

# 添加一个函数来获取详细的模块状态
def get_module_status() -> Dict[str, Any]:
    """
    获取模块的详细状态信息
    
    Returns:
        dict: 包含模块各组件状态的字典
    """
    return {
        "pyautogui_available": pyautogui is not None,
        "numpy_available": np is not None,
        "opencv_available": cv2 is not None,
        "tesseract_available": tesseract_available,
        "tesseract_path": getattr(pytesseract, 'tesseract_cmd', None) if pytesseract is not None else None,
        "can_screenshot": pyautogui is not None,
        "can_process_images": np is not None and cv2 is not None,
        "can_ocr": tesseract_available and np is not None,
        "recommended_installs": [
            "Tesseract OCR (https://github.com/UB-Mannheim/tesseract/wiki)" if not tesseract_available else None,
        ]
    }

# 尝试自动查找并配置Tesseract
try:
    if pytesseract is not None and not tesseract_available:
        tesseract_path = find_tesseract()
        if tesseract_path:
            set_tesseract_path(tesseract_path)
except Exception as e:
    logger.error(f"自动配置Tesseract时出错: {str(e)}")

# 初始化函数
def initialize() -> bool:
    """初始化CR模块
    
    Returns:
        是否初始化成功
    """
    global screen_recognizer
    screen_recognizer = ScreenContentRecognizer()
    success = screen_recognizer.tesseract_available
    
    if success:
        logger.info("CR模块初始化成功")
    else:
        logger.warning("CR模块初始化不完全，部分功能可能不可用")
    
    return success

# 如果直接运行此模块，执行简单测试
if __name__ == "__main__":
    print("=== CR模块测试 ===")
    
    if initialize():
        print("1. 获取屏幕信息:")
        info = get_screen_info()
        print(f"   屏幕尺寸: {info['screen_size']}")
        print(f"   鼠标位置: {info['mouse_position']}")
        print(f"   Tesseract可用: {info['tesseract_available']}")
        print(f"   OpenCV可用: {info['opencv_available']}")
        
        print("\n2. 提取屏幕文本:")
        text = extract_text_from_screen()
        print(f"   提取到的文本:\n{text}")
        
        print("\n3. 分析屏幕内容:")
        analysis = analyze_screen()
        print(f"   文本总数: {analysis['text_count']}")
        print(f"   文本区域数: {analysis['region_count']}")
    else:
        print("CR模块初始化失败，请安装Tesseract OCR")